import torch
import torch.nn as nn  # 导入PyTorch神经网络模块
import config  # 导入项目配置参数


class YoloV1Net(nn.Module):
    def __init__(self, S=config.S, B=config.B, C=config.C):
        """
        初始化YOLOv1网络模型
        :param S: 网格数量（S×S），对应config中的S
        :param B: 每个网格预测的边界框数，对应config中的B
        :param C: 目标类别数，对应config中的C
        """
        super(YoloV1Net, self).__init__()
        self.B = B  # 每个网格的边界框数B
        self.S = S  # 网格数量S
        self.C = C  # 类别数C
        self.depth = self.B * 5 + self.C  # 输出张量深度（2×5属性 + 20类别）
        self.Conv = self.make_Conv()  # 构建卷积层列表
        self.Classifier = self.make_Classifier()  # 构建全连接层列表
        self.conv_network = nn.Sequential (*self.Conv)  # 卷积网络（顺序执行各层）
        self.linear_network = nn.Sequential (*self.Classifier)  # 全连接网络（顺序执行各层）

    def forward(self, x):
        """
        前向传播计算（定义数据流经网络的路径）
        :param x: 输入图像数据（形状：(batch, 3, 448, 448)）
        :return: 网络输出（形状：(batch, S, S, depth)，默认(128,7,7,30)）
        """
        out = self.conv_network (x)  # 输入图像通过卷积网络提取特征
        out = self.linear_network (out)  # 卷积特征通过全连接网络预测边界框和类别
        # 调整输出形状：[batch, S*S*depth] -> [batch, S, S, depth]
        out = torch.reshape(out, (x.size(0), self.S, self.S, self.depth))
        # 说明：x.size(0)为批量大小，保持与输入一致
        return out

    def make_Conv(self):
        """
        构建YOLOv1的卷积网络（特征提取 backbone）
        输入图像尺寸：(448, 448, 3)（宽×高×通道）
        :return: 卷积层列表（按顺序执行）
        """
        # ---------------------------
        # 第1层：初始特征提取
        # ---------------------------
        layer1 = [
            nn.Conv2d(3, 64, 7, 2, 3),  # 输入3通道，输出64通道，7×7卷积核，步长2，填充3
            nn.LeakyReLU(negative_slope=0.1),  # 带负斜率的ReLU激活函数（防止死神经元）
            nn.MaxPool2d(2, 2)  # 2×2最大池化，步长2（尺寸减半）
        ] # 输出尺寸：(112, 112, 64)（宽高=448/2/2=112）

        # ---------------------------
        # 第2层：加深特征通道
        # ---------------------------
        layer2 = [
            nn.Conv2d(64, 192, 3, 1, 1),  # 输入64通道，输出192通道，3×3卷积核，步长1，填充1（保持尺寸）
            nn.LeakyReLU(negative_slope=0.1),
            nn.MaxPool2d(2, 2)  # 2×2最大池化（尺寸减半）
        ] # 输出尺寸：(56, 56, 192)（宽高=112/2=56）

        # ---------------------------
        # 第3层：1×3卷积交替（减少计算量）
        # ---------------------------
        layer3 = [
            nn.Conv2d(192, 128, 1, 1, 0),  # 1×1卷积（降维，减少计算量）
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(128, 256, 3, 1, 1),  # 3×3卷积（提取空间特征）
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(256, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.MaxPool2d(2, 2)
        ] # 输出尺寸：(28, 28, 512)（宽高=56/2=28）

        # ---------------------------
        # 第4层：重复1×3卷积块（加深网络）
        # ---------------------------
        layer4 = [
            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.MaxPool2d(2, 2)
        ] # 输出尺寸：(14, 14, 1024)（宽高=28/2=14）

        # ---------------------------
        # 第5层：调整尺寸至7×7（匹配S=7）
        # ---------------------------
        layer5 = [
            nn.Conv2d(1024, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            # 论文中的网络层（通过卷积下采样至7×7）
            nn.Conv2d(1024, 1024, 3, 2, 1),  # 3×3卷积，步长2，填充1（尺寸14→7）
            nn.LeakyReLU(negative_slope=0.1)

            # 自定义实现（使用池化下采样）
            # nn.MaxPool2d (2, 2)
        ] # 输出尺寸：(7, 7, 1024)（宽高=14/2=7，匹配S=7）

        # ---------------------------
        # 第6层：最终特征增强
        # ---------------------------
        layer6 = [
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1)
        ] # 输出尺寸：(7, 7, 1024)（保持7×7尺寸）

        return layer1 + layer2 + layer3 + layer4 + layer5 + layer6  # 拼接所有卷积层

    def make_Classifier(self):
        """
        构建YOLOv1的全连接网络（预测头）
        将卷积特征映射转换为边界框和类别预测
        :return: 全连接层列表（按顺序执行）
        """
        out_features = self.S * self.S * self.depth  # 输出特征数（S×S网格 × depth维）
        layer = [
            nn.Flatten(),  # 展平卷积特征（7×7×1024 → 7*7*1024=50176）
            nn.Linear(self.S * self.S * 1024, 4096),  # 全连接层（50176→4096）
            nn.Dropout(),  # Dropout层（防止过拟合，随机置零部分神经元）
            nn.LeakyReLU(negative_slope=0.1),  # 激活函数

            nn.Linear(4096, out_features)  # 输出层（4096→S×S×depth=7*7*30=1470）
        ]
        
        return layer  # 返回全连接层列表


if __name__ == '__main__' :
    net = YoloV1Net(config.S, config.B, config.C)
    print(net)